package cas.ihep.spark.rdd;

import cas.ihep.spark.context.JSparkIterationListener;
import cas.ihep.spark.loop.ClientIterativeRunnable;
import cas.ihep.spark.loop.IterationClient;
import cas.ihep.spark.loop.IterationData;
import cas.ihep.spark.loop.IterativeRunnable;
import cas.ihep.spark.util.*;
import cas.ihep.spark.util.function.Function2;
import cas.ihep.spark.util.function.Function4;
import org.apache.spark.Partition;
import org.apache.spark.SparkContext;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext$;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.rdd.RDD;
import scala.Tuple2;
import scala.collection.JavaConversions;
import scala.reflect.ClassManifestFactory;
import scala.reflect.ClassTag;

import java.io.Serializable;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.List;

public class JIterativePairRDD<K,V> extends JavaPairRDD<K,V> {

    private JIterativePairRDD(RDD<Tuple2<K, V>> rdd, ClassTag<K> kClassTag, ClassTag<V> vClassTag) {
        super(rdd, kClassTag, vClassTag);
    }

    public static <K,V> JIterativePairRDD<K,V> from(JavaPairRDD<K,V> rdd){
        if(rdd instanceof JIterativePairRDD){
            return (JIterativePairRDD<K, V>) rdd;
        }
        return new JIterativePairRDD<>(rdd.rdd(),rdd.kClassTag(),rdd.vClassTag());
    }

    public static <K,V> JIterativePairRDD<K,V> from(JavaRDD<Tuple2<K,V>> rdd){
        return new JIterativePairRDD<>(rdd.rdd(), JavaSparkContext$.MODULE$.<K>fakeClassTag(), JavaSparkContext$.MODULE$.<V>fakeClassTag());
    }

    public static <K,V> JIterativePairRDD<K,V> from(RDD<Tuple2<K,V>> rdd){
        return new JIterativePairRDD<>(rdd,JavaSparkContext$.MODULE$.<K>fakeClassTag(),JavaSparkContext$.MODULE$.<V>fakeClassTag());
    }

    public List<String> getPreferredLocations(Partition p){
        return JavaConversions.seqAsJavaList(rdd().getPreferredLocations(p));
    }

    public <Send extends Serializable, Receive extends Serializable>
    void iterate(IterativeRunnable<Send, Receive> iRun, IterationClientKVListener<K,V, Send, Receive> cRun,
                                                                                  Class<Send> pClass, Receive zero, Function2<Receive, Receive, Receive> agg){
        IterativeRDDUtil.iterate(rdd(),context(),getNumPartitions(),iRun,cRun,pClass,zero,agg);
    }

    public <Send extends Serializable, Receive extends Serializable>
    void iterate(IterativeRunnable<Send,Receive> iRun, IterationClientListener<Tuple2<K,V>,Send,Receive> cRun,
                 Class<Send> pClass,Receive zero,Function2<Receive,Receive,Receive> aggregateFunc){
        IterativeRDDUtil.iterate(rdd(),context(),getNumPartitions(),iRun,cRun,pClass,zero,aggregateFunc);
    }

    @Deprecated
    public <Send extends Serializable, Receive extends Serializable> void iterate(IterativeRunnable<Send, Receive> iRun,
                                                             Function4<K,V, Send, NamedResourceManager, Receive> cRun,
                                                             Class<Send> pClass, Receive zero, Function2<Receive, Receive, Receive> agg){
        final SparkContext sc=context();
        IterativeBroadcastDataKV<K,V, Send, Receive> param;
        try{
            param=new IterativeBroadcastDataKV<>(new InetSocketAddress(InetAddress.getLocalHost(),10000),pClass,zero,agg,cRun);
        }catch (UnknownHostException error){
            throw new RuntimeException("Should never happens",error);
        }
        param.properties=GlobalConfiguration.properties();
        sc.addSparkListener(new JSparkIterationListener<>(sc,iRun,getNumPartitions(),param.serverAddr));
        final Broadcast<IterativeBroadcastDataKV<K,V, Send, Receive>> bcast=sc.broadcast(param,ClassManifestFactory.classType(param.getClass()));
        rdd().foreach(new PairVoidFunction<K, V>() {
            @Override
            public void call(K key,V value) throws Exception {
                IterativeBroadcastDataKV<K,V, Send, Receive> param1=bcast.getValue();
                initExecutorProperties(param1.properties);
                try(NamedResourceManagerImpl closer=new NamedResourceManagerImpl()){
                    InetSocketAddress localAddr=new InetSocketAddress(InetAddress.getLocalHost(),8999);
                    @SuppressWarnings("unchecked")
                    IterationClient<Send, Receive> client1=IterationClient.getOrCreateSingleInstance(BaseBroadcastData.CLIENT_KEY,localAddr,
                            param1.serverAddr,(Class<Send>)param1.sendClassTag.runtimeClass());
                    //noinspection Convert2Lambda
                    client1.waitingForData(new ClientIterativeRunnable<Send, Receive>() {
                        @Override
                        public void run(IterationData<Send, Receive> iData) {
                            iData.output(param1.clientFunc.call(key,value,iData.input(),closer));
                        }
                    },param1.zero,param1.aggregateFunc);
                }
            }
        });
    }

    @Deprecated
    public <Send extends Serializable, Receive extends Serializable> void iterate(IterativeRunnable<Send, Receive> iterativeRunnable, final Function2<ManagedResourceKV<K,V>, Send, Receive> client, Class<Send> pClass, Receive zero,
                                                                                  Function2<Receive, Receive, Receive> aggregateFunc){
        final SparkContext sc=context();
        LoopBroadcastParam<Tuple2<K,V>, Send, Receive> param;
        try {
            param=LoopBroadcastParam.from(new InetSocketAddress(InetAddress.getLocalHost(), 10000),
                    zero,aggregateFunc,client, ClassManifestFactory.classType(pClass));
        }catch (UnknownHostException error){
            throw new RuntimeException("Should never happens",error);
        }
        param.props= GlobalConfiguration.properties();
        sc.addSparkListener(new JSparkIterationListener<>(sc, iterativeRunnable,getNumPartitions(),param.serverAddr));
        final Broadcast<LoopBroadcastParam<Tuple2<K,V>, Send, Receive>> globalConfigBcast=sc.broadcast(param, ClassManifestFactory.classType(param.getClass()));
        rdd().foreach(new PairVoidFunction<K, V>() {
            public void call(K key,V value)throws Exception{
                final LoopBroadcastParam<Tuple2<K, V>, Send, Receive> param1 = globalConfigBcast.getValue();
                initExecutorProperties(param1.props);
                try (DefaultKVResource<K, V> resource = new DefaultKVResource<>(key,value)) {
                    InetSocketAddress localAddr = new InetSocketAddress(InetAddress.getLocalHost(), 8999);
                    @SuppressWarnings("unchecked")
                    IterationClient<Send, Receive> client1 = IterationClient.getOrCreateSingleInstance(BaseBroadcastData.CLIENT_KEY,
                            localAddr, param1.serverAddr, (Class<Send>) param1.classTag.runtimeClass());
                    client1.waitingForData(lcr -> lcr.output(param1.client.call(resource, lcr.input())), param1.zero, param1.aggregate);
                }
            }
        });
    }

    private static abstract class PairVoidFunction<K,V> extends IterativeRDDUtil<Tuple2<K,V>>{

        public final void call(Tuple2<K,V> kvtp)throws Exception{
            call(kvtp._1(),kvtp._2());
        }
        public abstract void call(K key,V value)throws Exception;
    }

    private static class DefaultKVResource<K,V> extends DefaultManagedResource<Tuple2<K,V>> implements ManagedResourceKV<K,V>{

        private K k;
        private V v;
        DefaultKVResource(K kk,V vv){
            super(null);
            k=kk;
            v=vv;
        }

        public Tuple2<K,V> get(){return resource!=null?resource:(resource=new Tuple2<>(k,v));}

        public final K key(){
            return k;
        }

        public final V value(){
            return v;
        }
    }
}
